/**
 * @Auther: zhatiai
 * @Date:2024/10/9 15:06
 */

package iperf3

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os/exec"
)

type Iperf3Result struct {
	End *Iperf3EndStreams `json:"end"`
}

type Iperf3EndStreams struct {
	Streams []*Iperf3Stream `json:"streams"`
}

type Iperf3Stream struct {
	Sender   *Iperf3Bandwidth `json:"sender"`
	Receiver *Iperf3Bandwidth `json:"receiver"`
}

type Iperf3Bandwidth struct {
	BitsPerSecond float64 `json:"bits_per_second"`
}

type Iperf3 struct {
	serverHost string
	Host       string
	Port       int
}

func NewIperf3() *Iperf3 {
	return &Iperf3{}
}

func (i *Iperf3) Run(testTime int, isOutbound bool) (*Iperf3Result, error) {
	cmd := fmt.Sprintf("iperf3 -c %s -J", i.Host)
	if testTime == 0 {
		testTime = 3
	}

	cmd = fmt.Sprintf("%s -t %d", cmd, testTime)
	if isOutbound {
		cmd = fmt.Sprintf("%s %s", cmd, "-R")
	}

	command := exec.Command("sh", "-c", cmd)

	var outBuf, errBuf bytes.Buffer
	command.Stdout = &outBuf
	command.Stderr = &errBuf

	// 执行命令
	err := command.Run()
	if err != nil {
		// 捕捉错误输出
		return nil, fmt.Errorf("failed to execute iperf3 command: %v, stderr: %s, stdout: %s", err, errBuf.String(), outBuf.String())
	}

	// 解析 JSON 结果
	var result Iperf3Result
	err = json.Unmarshal(outBuf.Bytes(), &result)
	if err != nil {
		return nil, fmt.Errorf("failed to parse JSON output: %v, output: %s", err, outBuf.String())
	}

	return &result, nil
}

func (i *Iperf3) GetInboundBandwidth(host string) (*Iperf3Stream, error) {
	i.Host = host

	result, err := i.Run(3, false)
	if err != nil {
		return nil, err
	}

	if result.End != nil && result.End.Streams != nil && len(result.End.Streams) > 0 {
		return result.End.Streams[0], nil
	}

	return nil, fmt.Errorf("no bandwidth data found")
}

func (i *Iperf3) GetOutboundBandwidth(host string) (*Iperf3Stream, error) {
	i.Host = host

	result, err := i.Run(1, true)
	if err != nil {
		return nil, err
	}

	if result.End != nil && result.End.Streams != nil && len(result.End.Streams) > 0 {
		return result.End.Streams[0], nil
	}

	return nil, fmt.Errorf("no bandwidth data found")
}
